import json
import pprint
from io import BytesIO

from dotenv import load_dotenv
from sqlalchemy import false
from starlette.responses import JSONResponse

from src.chatAI import MyHistoryLoader, ZhiPuAI, AliAI
from src.custom.messages import build_system_message, build_human_message
from src.custom.tools import read_json, logger, set_logger, match_json_str
from src.file_parsers.FileManage import get_file_cache_dict, write_file_cache, get_file_cache_dict_by_
from src.file_parsers.PDFParsers import pdf_parser
from src.user import UserInfo

system_message = """
### 你是一个文件导读助手，尽可能的帮助用户通过文档里面的内容来回答，如果不存在则自己回答。
### 导读形式应该为：对每一块相关联的内容进行总结回答，尽可能压缩原有文本的内容，但要注意关键信息的提取。
### 你的任务是通过给出的文档参考内容回答问题，如果回答参考了文档里的内容，则需要标明参考的对应序号。如果没有参考则为空。
    一句话中的每一段文本可能对应不同的序号，这种情况下需要将他们在json中分开，但语义上连贯。
    每段话与其对应的序号颗粒度要细，尽可能准确定位到原文本信息。
### `<content></content>`标签括起来的是上传的文件内容。
### `<source_file_id></source_file_id>`标签内包裹的是<content>的来源id，表示这个内容的一个标识符。
### json格式说明：
{{
    "answers": [//是用来存放所有回答内容的，是一个数组，具体内容的格式，包含键如下：
        {{
            "text": "回答问题的文本，内容和序号尽量分细，不要聚合在一起",
            "order": [ //参考内容的对应序号
                {{"page_num": "页码", "index": "序号"}}
            ]
            "source_file_id": "回答内容的来源文件id，可以根据这个文件id追溯源文件"
        }}
    ]
}}
### 没有参考到<content>文件里的内容则order为[]，source_file_id为""。
### 例子1：
<source_file_id>
111111122222222222333333333444444444
</source_file_id>
<content>
[
    {{
        "page_num": 1,
        "index": 0,
        "content": "小红有一个弟弟叫小明。"
    }},
    {{
        "page_num": 1,
        "index": 1,
        "content": "小明的二弟叫小蓝。"
    }},
    {{
        "page_num": 2,
        "index": 2,
        "content": "他有两个弟弟,"
    }},
    {{
        "page_num": 2,
        "index": 3,
        "content": "三弟叫小黄"
    }},
    {{
        "page_num": 2,
        "index": 4,
        "content": "他们一家有7个人"
    }}
]
</content>
提问：“小明有几个弟弟，分别叫什么”
回答：
{{
    "answers":[
        {{"text": "根据您提供的内容可知：\n", "order": [], "source_file_id": ""}},
        {{"text": "小明家里面有7口人，", "order": [{{"page_num": 2, "index": 4}}], "source_file_id": "111111122222222222333333333444444444"}},
        {{"text": "其中小明有两个弟弟\n", "order": [{{"page_num": 2, "index": 2}}], "source_file_id": "111111122222222222333333333444444444"}},
        {{"text": "二弟叫小蓝，三弟叫小黄。\n", "order": [{{"page_num": 1, "index": 1}}, {{"page_num": 2, "index": 3}}], "source_file_id": "111111122222222222333333333444444444"}}
    ]
}}
### 参考信息：
<source_file_id>
{source_file_id}
</source_file_id>
<content>
{content}
</content>
"""

new_system_message = """
### 你是一个文件分析助手，尽可能的帮助用户通过文档里面的内容来回答，如果不存在则自己回答。
### 你的任务是通过给出的文档参考内容回答问题，如果回答参考了文档里的内容，则需要标明参考的对应序号。如果没有参考则为空。
    一句话中的每一段文本可能对应不同的序号，这种情况下需要将他们在json中分开，但语义上连贯。
    每段话与其对应的序号颗粒度要细，尽可能准确定位到原文本信息。
### `<content></content>`标签括起来的是上传的文件内容。
### `<source_file_id></source_file_id>`标签内包裹的是<content>的来源id，表示这个内容的一个标识符。
### json格式说明：
{
    "answers": [//是用来存放所有回答内容的，是一个数组，具体内容的格式，包含键如下：
        {
            "text": "回答问题的文本，内容和序号尽量分细，不要聚合在一起",
            "order": [ //参考内容的对应序号
                {"page_num": "页码", "index": "序号"}
            ]
            "source_file_id": "回答内容的来源文件id，可以根据这个文件id追溯源文件"
        }
    ]
}
### 没有参考到<content>文件里的内容则order为[]，source_file_id为""。
### 例子1：
<source_file_id>
111111122222222222333333333444444444
</source_file_id>
<content>
[
    {
        "page_num": 1,
        "index": 0,
        "content": "小红有一个弟弟叫小明。"
    },
    {
        "page_num": 1,
        "index": 1,
        "content": "小明的二弟叫小蓝。"
    },
    {
        "page_num": 2,
        "index": 2,
        "content": "他有两个弟弟,"
    },
    {
        "page_num": 2,
        "index": 3,
        "content": "三弟叫小黄"
    },
    {
        "page_num": 2,
        "index": 4,
        "content": "他们一家有7个人"
    }
]
</content>
提问：“小明有几个弟弟，分别叫什么”
回答：
{
    "answers":[
        {"text": "根据您提供的内容可知：\n", "order": [], "source_file_id": ""},
        {"text": "小明家里面有7口人，", "order": [{"page_num": 2, "index": 4}], "source_file_id": "111111122222222222333333333444444444"},
        {"text": "其中小明有两个弟弟\n", "order": [{"page_num": 2, "index": 2}], "source_file_id": "111111122222222222333333333444444444"},
        {"text": "二弟叫小蓝，三弟叫小黄。\n", "order": [{"page_num": 1, "index": 1}, {"page_num": 2, "index": 3}], "source_file_id": "111111122222222222333333333444444444"}
    ]
}
### 参考信息：
<source_file_id>
{source_file_id}
</source_file_id>
<content>
{content}
</content>
"""

lizi2 = """
### 例子2：
<source_file_id>
1234567890
</source_file_id>
<content>
[
    {
        "page_num": 1,
        "index": 0,
        "content": "这是一篇关于大熊猫的介绍。"
    },
    {
        "page_num": 1,
        "index": 1,
        "content": "大熊猫是中国特有的物种。"
    },
    {
        "page_num": 2,
        "index": 2,
        "content": "它们主要生活在中国四川、陕西和甘肃的山区。"
    }
]
</content>
提问：“简述一下这个文件的内容”
回答：
{
    "answers":[
        {"text": "该文件主要介绍了大熊猫的基本情况，", "order": [], "source_file_id": ""},
        {"text": "包括它是", "order": [{"page_num": 1, "index": 0}], "source_file_id": "1234567890"},
        {"text": "中国特有的物种", "order": [{"page_num": 1, "index": 1}], "source_file_id": "1234567890"},
        {"text": "以及它们的主要栖息地是在", "order": [], "source_file_id": ""},
        {"text": "中国四川、陕西和甘肃的山区", "order": [{"page_num": 2, "index": 2}], "source_file_id": "1234567890"}
    ]
}
"""

ref_content = new_system_message

def get_file_if_cached_otherwise_parse(file_path_or_bytes: str | bytes, **kwargs) -> tuple[dict, dict]:
    """
    如果有缓存的话获取缓存，否则解析PDF文件并写入缓存

    :param file_path_or_bytes: 需要获取内容的文件二进制或者文件路径
    :return: PDF信息，缓存字典
    """
    # 查看缓存
    cache_or_sha256 = get_file_cache_dict(file_path_or_bytes)
    if type(cache_or_sha256) is not str:
        logger.info(f"有文件缓存，正在读取缓存 . . .")
        # 有缓存读取
        block_info = read_json(cache_or_sha256['file_text_path'])
    else:
        logger.info(f"没有文件缓存，正在解析文件 . . .")
        # 无缓存解析
        if isinstance(file_path_or_bytes, bytes):
            file_like_obj = BytesIO(file_path_or_bytes)
            block_info = pdf_parser(file_like_obj)
        else:
            block_info = pdf_parser(file_path_or_bytes)
        logger.info(f"文件解析完毕！")
        logger.info(f"准备写缓存 . . .")
        # 写缓存
        write_file_cache(kwargs['file_name'], file_path_or_bytes, block_info, cache_or_sha256)
        cache_or_sha256 = get_file_cache_dict_by_(cache_or_sha256)
    logger.info(f"返回缓存信息")
    return block_info, cache_or_sha256

def get_extent_by_(current_page_n, pdf_info, extent: tuple = (0, 10)) -> list:
    """
    通过前端传递过来的当前页码，来获取基于当前页的前几页和后几页。
    比如current_page_n=4，默认就是取3~6页，包含第6页，一共4页
    :param pdf_info: 获取的解析后的pdf信息
    :param current_page_n: 当前预览的页码
    :param extent: 前几页个数，后几页个数
    :return: 前闭后开的区间，直接对数组使用，不需要额外加减
    """
    page_num = current_page_n
    u = page_num - extent[0]
    l = page_num + extent[1]
    upper = u if u >= 1 else 1
    lower = l if l <= len(pdf_info) else len(pdf_info)
    return pdf_info[upper-1: lower]

def get_context_from_file(pdf_info, current_page_n):
    split_block = get_extent_by_(current_page_n, pdf_info)
    text = []
    # 根据页码获取对应的chunk
    for b in split_block:
        for i, chunk in enumerate(b['chunks']):
            text.append({
                "page_num": b['page_num'],
                "index": i,
                "content": chunk['text'],
            })
    return text

def join_llm_response(llm_response: str, block_info):
    # 获取到AI响应，然后处理响应，拼接成传递给前端的消息格式
    # response_json_obj = match_json_str(llm_response)
    response_json_obj = json.loads(llm_response)
    if type(response_json_obj) is str:
        # 如果是字符串就说明有错误信息了
        return JSONResponse(status_code=500, content={"error": str(response_json_obj)})
    # 添加文本的坐标信息
    for answer in response_json_obj['answers']:
        coords = []
        if block_info is not None:
            for page_n_and_index in answer['order']:
                page_n = page_n_and_index['page_num']
                index = page_n_and_index['index']
                try:
                    arr = block_info[page_n]['chunks'][index]['coords']
                except IndexError as e:
                    logger.error(f"索引错误：{e}，页码：{page_n}，索引：{index}")
                    continue
                coords.append(arr)
        # 在原有信息增加属性
        answer["coords"] = coords
    return response_json_obj

def receive_messages_and_files(message: str, file: str | bytes | None, **kwargs):
    """

    :param message:
    :param file:
    :param kwargs: 必须要有文件名file_name、current_page_n
    :return:
    """
    load_dotenv('../.env')
    # TODO 应该思考一下用户和历史加载器的调用关系
    # 何时加载何时保存
    u = UserInfo("转换文件测试")
    u.load()

    # 前端要接收消息的格式
    page_info = {
        "file_path": "",        # 文件在服务器的位置，前端预览的时候需要在这个位置访问
        "block_info": [],       # 每一块文本信息，其中包括文本、引用的chunk序号、对应chunk的所有文本坐标
    }

    """
    这里的逻辑是保证第一次新用户一定有文件，之后才有可能没文件
    所以else逻辑里面，user对象一定会有一个文件
    """
    if file is not None:
        block_info, file_info_cache = get_file_if_cached_otherwise_parse(file, **kwargs)
        # 将要发送给大模型的文本数据
        global new_system_message
        let_system_message = new_system_message
        my_history = MyHistoryLoader(u, let_system_message)
        my_history.load()
        # 判断是否为第一次加载文件
        if not u.load():
            u.set_session_name(file_info_cache['file_hash'])
        else:
            # 如果有加载后的信息，则又上传了新的文件，表示开启新的会话
            my_history.start_new_session(file_info_cache['file_hash'])
    else:
        my_history = MyHistoryLoader(u)
        my_history.load()
        file_info_cache = get_file_cache_dict_by_(u.session_name)
        block_info = read_json(file_info_cache['file_text_path'])
    text = get_context_from_file(block_info, kwargs['current_page_n'])
    let_ref = system_message.format(content=text, source_file_id=file_info_cache['file_hash'])
    my_history.add_message(build_system_message(let_ref, True))
    # 给大模型发消息
    zhipuai = ZhiPuAI(my_history)
    # aliai = AliAI(my_history)
    llm_response = zhipuai.send_message(message)
    logger.debug(f"大模型返回的原始消息：{llm_response}")
    # 完成一次对话，保存历史
    my_history.save()

    response_json_obj = join_llm_response(llm_response, block_info)
    page_info["block_info"] = response_json_obj['answers']
    page_info["file_path"] = file_info_cache['file_path'][2:]       # 去除路径前多余的 “..”
    # 返回给前端
    return page_info

def receive_request_history():
    u = UserInfo("转换文件测试")
    u.load()
    my_history = MyHistoryLoader(u)
    my_history.load()
    file_info_cache_dict = get_file_cache_dict_by_(u.session_name)
    block_info = read_json(file_info_cache_dict['file_text_path'])

    history_messages = []
    for message in my_history.get_history():
        if message['role'] == 'user':
            history_messages.append({
                "role": "user",
                "content": message['content'],
            })
        elif message['role'] == 'assistant':
            llm_response = message['content']
            response_json_obj = join_llm_response(llm_response, block_info)
            history_messages.append({
                "role": "assistant",
                "content": response_json_obj['answers'],
            })
    return {
        "history": history_messages,
        "file_path": file_info_cache_dict['file_path'][2:]
    }




if __name__ == '__main__':
    set_logger()

    file_path = r"./sources/中期报告.pdf"
    # user_input = "首先帮我总结一下这个报告的内容，然后分部分来引领我导读。"
    user_input = "报告中的主人公遇到了什么问题，怎么解决的"
    # 模拟请求
    response = receive_messages_and_files(user_input, file_path)
    pprint.pprint(response)


